#include <cmath>
#include <cassert>
#include <iostream>
#include <fstream>

using namespace std;

// Compute State x(t) at time t
// Given: x0, u0, a, b (Diff.eq. x^prime = ax + bu)
// Time t is relative to 0
double compute_state(double a, double b, double x0, double u0, double t)
{
  /*
  double state_change = exp(a*t) * x0;
  double integral = ( exp(a*t) - 1 ) / a;
  double x_t = state_change + integral * b * u0;
  return x_t;
  */

  double bu0 = b*u0;
  double x_t;

  if (a == 0) {
    x_t = x0 + t * bu0;
  }
  else {
    double ax0bu0 = a*x0 + b*u0;
    x_t = exp(a*t) * ax0bu0 - bu0;
    x_t /= a;
  }

  cout << t << "\t" << x_t << endl;

  return x_t;
}

// plot the state in the time interval t0 to t, gran is the granularity
void plot_state(char *filename, double a, double b, double x0, double u0, double t0, double t, double gran)
{
  ofstream out(filename);

  double t_k = t0;
  double state;
  while (t_k < t) {
    state = compute_state(a, b, x0, u0, t - t_k);
    
    out << t_k << "\t" << state << endl;

    t_k += gran;
    if (t_k > t)
      t_k = t;
  }
}

// integral of the squared state, T is the length of the integration
double integral_of_quadratic_state(double a, double b, double x0, double u0, double T)
{
  double bu0 = b*u0;
  double sum;

  if (a == 0) {
    sum = bu0*bu0 * T*T*T / 3;
    sum += x0 * bu0 * T*T;
    sum += x0*x0 * T;
  }
  else {
    double ax0bu0 = a*x0 + b*u0;
    double coeff1 = ax0bu0 * ax0bu0 / (2*a);
    double coeff2 = 2 * ax0bu0 * bu0 / a;
    sum = coeff1 * ( exp(2*a*T) - 1 );
    sum += coeff2 * ( 1 - exp(a*T) );
    sum += bu0 * bu0 * T;
    sum /= a*a;
  }

  return sum;
}

// same as above, t_i is the start time of the next control update, k is the control gain
double state_cost(double a, double b, double x0, double u0, double T, 
		  double t_nextupdate, double k)
{
  assert(t_nextupdate >= 0 && t_nextupdate <= T);

  // compute cost from 0 to t_nextupdate
  double cost1 = integral_of_quadratic_state(a, b, x0, u0, t_nextupdate);

  // compute cost from t_nextupdate to T
  double x1 = compute_state(a, b, x0, u0, t_nextupdate);
  double u1 = k*x1;
  double cost2 = integral_of_quadratic_state(a, b, x1, u1, T-t_nextupdate);
  
  return cost1+cost2;
}

void plot_state_cost(char *filename, double a, double b, double x0, double u0, 
		     double k, double t0, double t, double gran)
{
  ofstream out(filename);

  double t_k;
  double cost;

  t_k = t0;
  while (t_k < t) {
    cost = state_cost(a, b, x0, u0, t-t0, t_k-t0, k);
    
    out << t_k << "\t" << cost << endl;

    t_k += gran;
    if (t_k > t)
      t_k = t;
  }
}

void plot_normalized_state_cost(char *filename, double a, double b, double x0, double u0, 
				double k, double t0, double t, double gran)
{
  ofstream out(filename);

  double min_cost = 100000;
  double max_cost = -100000;

  double t_k;
  double cost;

  // find min and max costs
  t_k = t0;
  while (t_k < t) {
    cost = state_cost(a, b, x0, u0, t-t0, t_k-t0, k);
    
    if (cost < min_cost)
      min_cost = cost;
    
    if (cost > max_cost)
      max_cost = cost;

    t_k += gran;
    if (t_k > t)
      t_k = t;
  }

  // normalized costs
  t_k = t0;
  while (t_k < t) {
    cost = state_cost(a, b, x0, u0, t-t0, t_k-t0, k);
    cost = (cost - min_cost) / (max_cost - min_cost);

    out << t_k << "\t" << cost << endl;

    t_k += gran;
    if (t_k > t)
      t_k = t;
  }
}

void plot_cpu_cost(char *filename, double t0, double t, double gran)
{
  ofstream out(filename);

  double t_k;
  double cost;

  t_k = t0;
  while (t_k < t) {
    cost = (t-t_k) / (t-t0);
    
    out << t_k << "\t" << cost << endl;

    t_k += gran;
    if (t_k > t)
      t_k = t;
  }
}

// rho is the trade-off parameter between control cost and cpu cost
void plot_task_cost(char *filename, double a, double b, double x0, double u0, 
		    double k, double t0, double t, double gran, double rho)
{
  ofstream out(filename);

  double min_cost = 100000;
  double max_cost = -100000;

  double t_k;
  double cost;
  double cpucost;
  double taskcost;

  // find min and max costs
  t_k = t0;
  while (t_k < t) {
    cost = state_cost(a, b, x0, u0, t-t0, t_k-t0, k);
    
    if (cost < min_cost)
      min_cost = cost;
    
    if (cost > max_cost)
      max_cost = cost;

    t_k += gran;
    if (t_k > t)
      t_k = t;
  }

  // normalized costs
  t_k = t0;
  while (t_k < t) {
    cost = state_cost(a, b, x0, u0, t-t0, t_k-t0, k);
    cost = (cost - min_cost) / (max_cost - min_cost);

    cpucost = (t-t_k) / (t-t0);
    taskcost = cost + rho*cpucost;

    out << t_k << "\t" << taskcost << endl;

    t_k += gran;
    if (t_k > t)
      t_k = t;
  }
}

int main()
{
  double a = 0.5;
  double b = 1;
  double k = -1;
  double x0 = 1;
  double u0 = -1;
  double T = 3;
  double gran = 0.05;
  double rho = 1.5;

  double t0 = 10;

  plot_normalized_state_cost("statecost.dat", a, b, x0, u0, k, t0, t0+T, gran);
  plot_cpu_cost("cpucost.dat", t0, t0+T, gran);
  //  plot_task_cost("taskcost_rho0_5.dat", a, b, x0, u0, k, t0, t0+T, gran, 0.5);
  plot_task_cost("taskcost_rho1.dat", a, b, x0, u0, k, t0, t0+T, gran, 1);
  //  plot_task_cost("taskcost_rho1_5.dat", a, b, x0, u0, k, t0, t0+T, gran, 1.5);
  plot_task_cost("taskcost_rho2.dat", a, b, x0, u0, k, t0, t0+T, gran, 2);
  //  plot_task_cost("taskcost_rhoinf.dat", a, b, x0, u0, k, t0, t0+T, gran, 10);

  //  double cost = integral_of_quadratic_state(a, b, x0, u0, T);
  //  cout << cost << endl;

  return 0;
}
